#!/usr/bin/bash python
# coding=utf-8

from typing import Union, List
import psycopg2
import psycopg2.extras


class DBBase:
    def __init__(self, database: str, user: str, password: str, port: int, host: str):
        self.conn = psycopg2.connect(database=database, user=user, password=password, port=port, host=host)
        self.conn.autocommit = False
        self.cur = self.conn.cursor()

    def close(self):
        self.conn.close()

    def execute(self, sql: str, params: dict = None):
        try:
            if params:
                self.cur.execute(sql, params)
            else:
                self.cur.execute(sql)
            self.conn.commit()
        except Exception as e:
            self.conn.rollback()
            raise e

    def batch_execute(self, sqls: list, params: dict = None):
        try:
            for sql in sqls:
                if params:
                    self.cur.execute(sql, params)
                else:
                    self.cur.execute(sql)
            self.conn.commit()
        except Exception as e:
            self.cur.rollback()
            raise e

    def execute_by_execute_values(self, sql: str, data: Union[list, set], template: str = None, page_size: int = 100,
                                  fetch: bool = False):
        try:
            psycopg2.extras.execute_values(self.cur, sql, data, template=template, page_size=page_size, fetch=fetch)
            self.cur.commit()
        except Exception as e:
            self.cur.rollback()
            raise e

    def execute_by_execute_batch(self, sql: str, data: List[dict], page_size: int = 100):
        try:
            psycopg2.extras.execute_batch(self.cur, sql, data, page_size=page_size)
            self.cur.commit()
        except Exception as e:
            self.cur.rollback()
            raise e


if __name__ == '__main__':
    db = DBBase(database='test_db', user='admin', password='123456', port=5432, host='192.168.100.4')
    sql1 = 'create table p_report_1(id integer not null primary key,name text not null unique,style text, project text);'
    sql2 = 'create table p_report_2(id integer not null primary key,name text not null unique,style text, project text);'
    sql3 = 'create table p_report_3(id integer not null primary key,name text not null unique,style text, project text);'
    sqls = [sql1, sql2, sql3]
    db.batch_execute(sqls)
    db.close()
